""" Get the rho parameters that relate mri to value

"""

import pandas as pd
import statsmodels.formula.api as smf

#%% IMPORT DATA -------------------------------------------------------------------------
df_contracts = pd.read_csv('./data_py/processed/contracts_final.csv', index_col=[0])

#%% GET COEFFICIENTS --------------------------------------------------------------------
df_contracts['mri_2'] = df_contracts['mri'] * df_contracts['mri']
df_contracts['mri_3'] = df_contracts['mri'] * df_contracts['mri'] * df_contracts['mri']

formula = 'value ~ mri + mri_2 + mri_3'

reg = smf.ols(
    formula=formula,
    data=df_contracts
).fit(cov_type='HC0')

coefs_data = reg.params.T.rename(index={
    'Intercept': 'rho_0',
    'mri': 'rho_1',
    'mri_2': 'rho_2',
    'mri_3': 'rho_3'
})
ci = reg.conf_int(alpha=0.05, cols=None)
ci = ci.rename(index={
    'Intercept': 'rho_0',
    'mri': 'rho_1',
    'mri_2': 'rho_2',
    'mri_3': 'rho_3'
})

pd.Series(coefs_data).to_csv('./models/smm_input/rho.csv')
ci[0].to_csv('./models/smm_input/rho_lower.csv')
ci[1].to_csv('./models/smm_input/rho_upper.csv')
